#include <pr/NearestMeanClassifier.h>
#include <pr/Pattern.h>
#include <gtest/gtest.h>

using namespace pr;

TEST(TestNearestMeanClassifier, testCtor00)
{
    NearestMeanClassifier c(10, 20);

    ASSERT_EQ(10, c.classCount());
    ASSERT_EQ(20, c.featureCount());
}

TEST(TestNearestMeanClassifier, test1class)
{
    NearestMeanClassifier c(1, 4);

    Pattern trainingPattern0(4);
    trainingPattern0[0] = 0;
    trainingPattern0[1] = 0.5;
    trainingPattern0[2] = 0.7;
    trainingPattern0[3] = 0.4;

    c.train(trainingPattern0, 0);

    Pattern testPattern(4);
    testPattern[0] = 0;
    testPattern[1] = 0.1;
    testPattern[2] = 0.2;
    testPattern[3] = 0.3;

    ASSERT_TRUE(c.test(trainingPattern0, 0));
    ASSERT_TRUE(c.test(testPattern, 0));
}

TEST(TestNearestMeanClassifier, test2classes)
{
    NearestMeanClassifier c(2, 4);

    Pattern trainingPattern0(4);
    trainingPattern0[0] = 0;
    trainingPattern0[1] = 0.5;
    trainingPattern0[2] = 0.7;
    trainingPattern0[3] = 0.4;

    Pattern trainingPattern1(4);
    trainingPattern1[0] = 0;
    trainingPattern1[1] = 0.2;
    trainingPattern1[2] = 0.3;
    trainingPattern1[3] = 0.4;

    c.train(trainingPattern0, 0);
    c.train(trainingPattern1, 1);

    Pattern testPattern(4);
    testPattern[0] = 0;
    testPattern[1] = 0.1;
    testPattern[2] = 0.2;
    testPattern[3] = 0.3;

    ASSERT_TRUE(c.test(trainingPattern0, 0));
    ASSERT_TRUE(c.test(trainingPattern1, 1));
    ASSERT_TRUE(c.test(testPattern, 1));
}
